# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: chunhui

import json
import os
import pandas as pd
import shutil
import subprocess
import tempfile
import time
from absl import app, flags
from coin.exchange.okex_futures_v3.kr_rest.native_private_client import \
    OkexFuturesNativePrivateClient

FLAGS = flags.FLAGS


class OkexSwapClient(OkexFuturesNativePrivateClient):
  def __init__(self, key_file):
    super().__init__(key_file=key_file)
    with open('data/coin2/product_info/Futures.Okex.v5-swap.json') as f:
      product_info = json.load(f)
    self.mea = product_info['mea']
    self.product_info = product_info['product_infos']

  def convert_native_symbol(self, symbol):
    for info in self.product_info:
      if info['symbol'] == symbol:
        return info['native_symbol']
    raise ValueError(f'Unrecognized symbol {symbol}')

  def _set_leverage(self, symbol, leverage):
    method = 'POST'
    path = '/api/v5/account/set-position-mode'
    params = {
       "posMode":"net_mode"
    }
    update = self._query(method, path, params)

    path = '/api/v5/account/set-leverage'
    native_symbol = self.convert_native_symbol(symbol)

    params = {
        'instId': native_symbol,
        'lever': leverage,
        'mgnMode': 'cross',
        'posSide': 'net'
    }
    update = self._query(method, path, params)
    return update

  def set_leverage(self, symbols, leverage):
    if isinstance(symbols, str):
      symbols = symbols.split(',')
    for i, symbol in enumerate(symbols):
      if (i + 1) % 5 == 0:
        time.sleep(2)
      result = self._set_leverage(symbol, leverage)
      print(result)

  def _get_pos_limit(self, symbol):
    method = 'GET'
    path = '/api/v5/account/max-size'
    native_symbol = self.convert_native_symbol(symbol)
    params = {
        "instId": native_symbol,
        "tdMode": "cross",
        "ccy": "USDT"
    }
    update = self._query(method, path, params)
    return update

  def _get_trade_limit(self, symbol):
    method = 'GET'
    path = '/api/v5/account/max-avail-size'
    native_symbol = self.convert_native_symbol(symbol)
    params = {
        "instId": native_symbol,
        "tdMode": "cross",
        "ccy": "USDT"
    }
    update = self._query(method, path, params)
    return update

  def _get_leverage(self, symbol):
    method = 'GET'
    path = '/api/v5/account/leverage-info/'
    native_symbol = self.convert_native_symbol(symbol)
    params = {
        "instId": native_symbol,
        "mgnMode": "cross"
    }
    update = self._query(method, path, params)
    return update

  def get_leverage(self, symbols):
    if isinstance(symbols, str):
      symbols = symbols.split(',')
    data = []
    for i, symbol in enumerate(symbols):
      if (i + 1) % 19 == 0:
        time.sleep(2)
      leverage = self._get_leverage(symbol)
      pos_limit = self._get_pos_limit(symbol)
      trade_limit = self._get_trade_limit(symbol)
      data.extend(leverage.as_json()['msg']['data'])
      data[-1].update(pos_limit.as_json()['msg']['data'][0])
      data[-1].update(trade_limit.as_json()['msg']['data'][0])
    df = pd.DataFrame(data)
    return df


def init_flags():
  flags.DEFINE_string('strategy', '', 'strategy_name')
  flags.DEFINE_string('account', '', 'account name')
  flags.DEFINE_string('symbol', 'BTC-USDT.PERPETUAL', 'usdt amount')
  flags.DEFINE_integer('leverage', 10, 'upd leverage')
  flags.DEFINE_bool('query_only', True, 'check leverage')


def get_okex_account(strategy):
  if os.path.exists(f'python/coin_deploy/basis_strat2/driver/{strategy}.json'):
    driver = json.load(open(f'python/coin_deploy/basis_strat2/driver/{strategy}.json'))
    for _, config in driver['order'].items():
      if config['mea'] == 'Futures.Okex.v5-swap':
        return config['connection_config']['key_filepath'].split('/')[3], config['products']
  raise ValueError(f"look up account for {strategy} failed")


def main(_):
  pd.set_option('display.max_rows', None)
  account = flags.FLAGS.account
  strat_symbols = [
        "NYM-USDT.PERPETUAL",
  ]

  tmp_dir = tempfile.mkdtemp()
  try:
    key_file = os.path.join(tmp_dir, 'key.json')
    cmd = f"""./python/coin_deploy/fab.sh \
-f python/coin_deploy/fabhelper/fabutil.py \
-P upload_coin_key_high:vault_token_file="./python/coin_deploy/vault/coin_key_high.json",\
key_name="trade/{account}/OKEx/trade_v5_key.json",\
remote_key_path="{key_file}",\
run_local=True
"""
    _ = subprocess.check_output(cmd, shell=True)
    msc = OkexSwapClient(key_file=key_file)

    result = msc.get_leverage(strat_symbols)
    print("\n#####################")
    print("-- Before Adjust --")
    print(result)
    print("######################\n")

    if not flags.FLAGS.query_only:
      strat_symbols = ",".join(strat_symbols)
      msc.set_leverage(strat_symbols, FLAGS.leverage)
      print("\n#####################")
      print("-- After Adjust --")
      print(msc.get_leverage(strat_symbols))
      print("######################\n")
  except Exception:
    raise
  finally:
    shutil.rmtree(tmp_dir)


'''
./pyrunner cc/appcoin2/strategy/basis_strat/support/adjust_leverage/okex_leverage.py
--strategy=mmrev_okex_perpetual_1
--query_only=0
--leverage=10
--symbol=BTC-USDT.PERPETUAL,ETH-USDT.PERPETUAL
'''

if __name__ == '__main__':
  init_flags()
  app.run(main)
